#!/usr/bin/env python
#
# Copyright (C) 2011 Clement Lorteau <clem@lorteau.fr>	
#		
# This file is part of QuickScreenShots.		
# 	
# This program is free software; you can redistribute it and/or modify	
# it under the terms of the GNU General Public License as published by		
# the Free Software Foundation; version 3.	
# 	
# This program is distributed in the hope that it will be useful,		
# but WITHOUT ANY WARRANTY; without even the implied warranty of		
# MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE.  See the		
# GNU General Public License for more details.		
# 		
# You should have received a copy of the GNU General Public License		
# along with this program; if not, write to the Free Software		
# Foundation, Inc., 51 Franklin St, Fifth Floor, Boston, MA 02110-1301, USA.

import wx
import os
import timer
import pyHook
from tempfile import NamedTemporaryFile
from PIL import Image
import win32gui,  win32ui,  win32con, win32api

class snapEvent(wx.CommandEvent):
    def __init__(self, event_type, id):
        wx.CommandEvent.__init__(self,event_type, id)

class escEvent(wx.CommandEvent):
    def __init__(self, event_type, id):
        wx.CommandEvent.__init__(self,event_type, id)

class magnifierFrame(wx.Frame):
    def __init__(self, evthdl, include_cursor = True, *args, **kwds):
        """
        A frame that magnifies the part of the screen around the pointer to ease
        selection of a region to make a screenshot of
        
        evthdl:     the EventHandler object that will handle the events fired by the frame.
                   
        Such events are:
        - magnifierFrame.EVT_SNAP_EVENT:    screenshot was taken and the frame destroyed; the
                                            bitmap is in magnifierFrame.screenshotbmp and the
                                            PIL Image is in magnifierFrame.screenshotim
        - magnifierFrame.EVT_ESC_EVENT:     user pressed the escape key; screenshot was cancelled
                                            and the frame destroyed
        """
        self.deskWnd = None
        self.snapPoint1 = None
        self.snapPoint2 = None
        self.screenshotim = None
        self.screenshotbmp = None
        self.evthdl = evthdl
        self.lastX = self.lastY = self.lastC = -1
        
        self.snap_event = wx.NewEventType()
        self.EVT_SNAP_EVENT = wx.PyEventBinder(self.snap_event, 1)
        self.esc_event = wx.NewEventType()
        self.EVT_ESC_EVENT = wx.PyEventBinder(self.esc_event, 1)
        
        self.freezeDesktop(include_cursor)
		
        # begin wxGlade: magnifierFrame.__init__
        kwds["style"] = wx.STAY_ON_TOP
        wx.Frame.__init__(self, *args, **kwds)
        self.panel_1 = wx.Panel(self, -1, style=wx.SUNKEN_BORDER|wx.TAB_TRAVERSAL)
        self.magnifyBitmap = wx.StaticBitmap(self.panel_1, -1, wx.NullBitmap, style=wx.SUNKEN_BORDER)
        self.label_1 = wx.StaticText(self, -1, "X:")
        self.label_2 = wx.StaticText(self, -1, "Y:")
        self.label_3 = wx.StaticText(self, -1, "Color:")
        self.label_4 = wx.StaticText(self, -1, "Use the arrow keys or the mouse to move the cursor. Use the <enter> key or left click to set a corner.", style=wx.ST_NO_AUTORESIZE)
        self.xLabel = wx.StaticText(self, -1, "label_4")
        self.yLabel = wx.StaticText(self, -1, "label_4")
        self.cLabel = wx.StaticText(self, -1, "label_4")

        self.__set_properties()
        self.__do_layout()
        # end wxGlade

        self.hm = pyHook.HookManager()
        self.hm.KeyDown = self.onKeyboard
        self.hm.MouseLeftDown = self.onMouseLeft
        self.hm.HookKeyboard()
        self.hm.HookMouse()
        self.timer = timer.set_timer(20, self.onTimer) #every 20 milliseconds

    def __set_properties(self):
        # begin wxGlade: magnifierFrame.__set_properties
        self.SetTitle("frame_1")
        self.SetBackgroundColour(wx.SystemSettings_GetColour(wx.SYS_COLOUR_3DFACE))
        self.magnifyBitmap.SetMinSize((256, 256))
        # end wxGlade
        self.label_4.Wrap(256)
        self.magnifyBitmap.SetDoubleBuffered(True)

    def __do_layout(self):
        # begin wxGlade: magnifierFrame.__do_layout
        sizer_2 = wx.BoxSizer(wx.VERTICAL)
        sizer_3 = wx.BoxSizer(wx.HORIZONTAL)
        sizer_5 = wx.BoxSizer(wx.VERTICAL)
        sizer_4 = wx.BoxSizer(wx.VERTICAL)
        sizer_1 = wx.BoxSizer(wx.HORIZONTAL)
        sizer_2.Add(self.label_4, 0, wx.ALL, 2)
        sizer_1.Add(self.magnifyBitmap, 0, 0, 0)
        self.panel_1.SetSizer(sizer_1)
        sizer_2.Add(self.panel_1, 1, 0, 0)
        sizer_4.Add(self.label_1, 0, wx.TOP|wx.BOTTOM, 2)
        sizer_4.Add(self.label_2, 0, wx.TOP|wx.BOTTOM, 2)
        sizer_4.Add(self.label_3, 0, wx.TOP|wx.BOTTOM, 2)
        sizer_3.Add(sizer_4, 1, wx.EXPAND, 3)
        sizer_5.Add(self.xLabel, 0, wx.TOP|wx.BOTTOM, 2)
        sizer_5.Add(self.yLabel, 0, wx.TOP|wx.BOTTOM, 2)
        sizer_5.Add(self.cLabel, 0, wx.TOP|wx.BOTTOM, 2)
        sizer_3.Add(sizer_5, 1, wx.EXPAND, 0)
        sizer_2.Add(sizer_3, 0, wx.ALL|wx.EXPAND, 3)
        self.SetSizer(sizer_2)
        sizer_2.Fit(self)
        self.Layout()
        # end wxGlade
        
    def getWXDCfromHandle(self, hwnd):
        window = wx.Frame(None, -1, '')
        window.AssociateHandle(hwnd)
        return wx.WindowDC(window)
		
    def freezeDesktop(self, getcursor = True):
        """
        Simulates freezing the desktop by creating a new top level window filled with
        the desktop contents
        """
        self.getDesktopDC(getcursor)
        wc = win32gui.WNDCLASS()
        wc.lpszClassName = 'man_is_win32_complicated'
        wc.hCursor = win32gui.LoadCursor(0, win32con.IDC_CROSS)
        self.class_atom = win32gui.RegisterClass(wc)
        deskWnd = win32gui.CreateWindow(wc.lpszClassName, '',
                                        win32con.WS_VISIBLE | win32con.WS_POPUPWINDOW,
                                        0, 0,
                                        self.screenwidth, self.screenheight,
                                        0, 0, 0,
                                        None)
        deskDC = win32ui.CreateDCFromHandle(win32gui.GetWindowDC(deskWnd))
        deskDC.BitBlt((0, 0), (self.screenwidth, self.screenheight),
                      self.desktopDC, (0, 0), win32con.SRCCOPY)
        self.deskWnd = deskWnd
            
    def getDesktopDC(self, getcursor = True):
        hwnd = win32gui.GetDesktopWindow()
        # get complete virtual screen including all monitors
        SM_XVIRTUALSCREEN = 76
        SM_YVIRTUALSCREEN = 77
        SM_CXVIRTUALSCREEN = 78
        SM_CYVIRTUALSCREEN = 79
        w = vscreenwidth = win32api.GetSystemMetrics(SM_CXVIRTUALSCREEN)
        h = vscreenheigth = win32api.GetSystemMetrics(SM_CYVIRTUALSCREEN)
        l = vscreenx = win32api.GetSystemMetrics(SM_XVIRTUALSCREEN)
        t = vscreeny = win32api.GetSystemMetrics(SM_YVIRTUALSCREEN)
        r = l + w
        b = t + h
        self.screenwidth = w
        self.screenheight = h
         
        hwndDC = win32gui.GetWindowDC(hwnd)
        mfcDC  = win32ui.CreateDCFromHandle(hwndDC)
        saveDC = mfcDC.CreateCompatibleDC()

        saveBitMap = win32ui.CreateBitmap()
        saveBitMap.CreateCompatibleBitmap(mfcDC, w, h)
        saveDC.SelectObject(saveBitMap)
        saveDC.BitBlt((0, 0), (w, h),  mfcDC,  (l, t),  win32con.SRCCOPY)

        #add cursor
        if getcursor:
            curFlags, curH, (curX, curY) = win32gui.GetCursorInfo()
            curWidth = win32api.GetSystemMetrics(win32con.SM_CXCURSOR)
            curHeight = win32api.GetSystemMetrics(win32con.SM_CYCURSOR)
            relCurX = curX - l - curWidth/4
            relCurY = curY - t - curHeight/4
            saveDC.DrawIcon((relCurX, relCurY), curH)
    
        self.desktopDC = saveDC
        
    def drawSnapPoint(self, point):
        deskDC = self.getWXDCfromHandle(self.deskWnd)
        deskDC.SetLogicalFunction(wx.INVERT)
        deskDC.SetBrush(wx.Brush(wx.Colour(0, 0, 0)))
        for cx in range (0, self.screenwidth):
            deskDC.DrawPoint(cx, point[1])
        for cy in range (0, self.screenheight):
            if cy != self.screenheight/2:
                deskDC.DrawPoint(point[0], cy)
                
    def snap(self):
        #determine square of region to be captured and upper left coordinates
        x1 = self.snapPoint1[0]
        y1 = self.snapPoint1[1]
        x2 = self.snapPoint2[0]
        y2 = self.snapPoint2[1]
        if x1 < x2:
            rectW = x2 - x1
            if y1 < y2: #first point is upper left
                rectH = y2 - y1
                topLeft = self.snapPoint1
            elif y2 < y1: #first point is bottom left
                rectH = y1 - y2
                topLeft = (x1, y2)
            else: #y1 = y2
                pass
        elif x1 > x2:
            rectW = x1 - x2
            if y1 < y2: #first point is upper right
                rectH = y2 - y1
                topLeft = (x2, y1)
            elif y1 > y2: #first point is bottom right
                rectH = y1 - y2
                topLeft = self.snapPoint2
            else: #y1 = y2
                pass
        else: #x1 = x2
            pass
                
        """rectW = max(self.snapPoint1[0], self.snapPoint2[0]) - min(self.snapPoint1[0], self.snapPoint2[0])
        rectH = max(self.snapPoint1[1], self.snapPoint2[1]) - min(self.snapPoint1[1], self.snapPoint2[1])
        if (self.snapPoint1[0] <= self.snapPoint2[0]) and (self.snapPoint1[1] <= self.snapPoint2[1]):
            topLeft = self.snapPoint1
        else:
            topLeft = self.snapPoint2"""
        
        #copy the determined square
        saveDC = self.desktopDC.CreateCompatibleDC()
        saveBitMap = win32ui.CreateBitmap()
        saveBitMap.CreateCompatibleBitmap(self.desktopDC, rectW, rectH)
        saveDC.SelectObject(saveBitMap)
        saveDC.BitBlt((0, 0), (rectW, rectH), self.desktopDC, topLeft, win32con.SRCCOPY)

        #load into PIL image
        """http://stackoverflow.com/questions/4199497/image-frombuffer-with-16-bit-image-data"""
        """bmpinfo = saveBitMap.GetInfo()
        bmpstr = saveBitMap.GetBitmapBits(True)
        try:
            im = Image.frombuffer(
                'RGB',
                (bmpinfo['bmWidth'], bmpinfo['bmHeight']),
                bmpstr, 'raw', 'BGRX', 0, 1)
        except ValueError, e: #32 bits didn't work, trying 16 bits
            im = Image.frombuffer(
                'RGB',
                (bmpinfo['bmWidth'], bmpinfo['bmHeight']),
                bmpstr, 'raw', 'BGR;16', 0, 1)"""
        """Avoid dealing with multiple color depth (the above doesn't work with 24
        bits and was not tested with <16 bits), write the bitmap to disk then load
        it from there"""
        tempfile = NamedTemporaryFile(prefix='qs_', delete=False)
        tempfile.close()
        saveBitMap.SaveBitmapFile(saveDC, tempfile.name)
        im = Image.open(tempfile.name)
        im.load()
        os.remove(tempfile.name)

        win32gui.DeleteObject(saveBitMap.GetHandle())
        saveDC.DeleteDC()

        #save					  
        self.screenshotim = im
        self.screenshotbmp = None
        
        #fire event
        snap_evt_obj = snapEvent(self.snap_event, 0)
        self.evthdl.ProcessEvent(snap_evt_obj)
        
    def Destroy(self):
        self.hm.UnhookMouse()
        self.hm.UnhookKeyboard()
        timer.kill_timer(self.timer)
        win32ui.CreateWindowFromHandle(self.deskWnd).DestroyWindow()
        win32gui.UnregisterClass(self.class_atom, None)
        wx.Frame.Destroy(self)
        
    def onKeyboard(self, event):
        if event.KeyID == 27:
            return self.onEsc()
        if event.KeyID == 38:
            return self.onUp()
        if event.KeyID == 40:
            return self.onDown()
        if event.KeyID == 37:
            return self.onLeft()
        if event.KeyID == 39:
            return self.onRight()
        if event.KeyID == 13:
            return self.onEnter()
    
    def onUp(self):
        curFlags, curH, (curX, curY) = win32gui.GetCursorInfo()
        win32api.SetCursorPos((curX, curY - 1))
        return False
        
    def onDown(self):
        curFlags, curH, (curX, curY) = win32gui.GetCursorInfo()
        win32api.SetCursorPos((curX, curY + 1))
        return False
        
    def onLeft(self):
        curFlags, curH, (curX, curY) = win32gui.GetCursorInfo()
        win32api.SetCursorPos((curX - 1, curY))
        return False
    
    def onRight(self):
        curFlags, curH, (curX, curY) = win32gui.GetCursorInfo()
        win32api.SetCursorPos((curX + 1, curY))
        return False
    
    def onEnter(self):
        curFlags, curH, (curX, curY) = win32gui.GetCursorInfo()
        win32api.mouse_event(win32con.MOUSEEVENTF_LEFTDOWN, curX, curY, 0, 0)
        win32api.mouse_event(win32con.MOUSEEVENTF_LEFTUP, curX, curY, 0, 0)
        return False
            
    def onMouseLeft(self, event):
        if not self.snapPoint1:
            self.snapPoint1 = event.Position
            self.drawSnapPoint(event.Position)
            return False
        if not self.snapPoint2:
            self.snapPoint2 = event.Position
            self.drawSnapPoint(event.Position)
            self.snap()
            return False
    
    def onEsc(self):
        #fire event
        esc_evt_obj = escEvent(self.esc_event, 0)
        self.evthdl.ProcessEvent(esc_evt_obj)
        return False
    
    def onTimer(self, timer_id, time):
        (x, y) = win32api.GetCursorPos()
        c = self.desktopDC.GetPixel(x, y)
        try:
            if x != self.lastX:
                self.xLabel.SetLabel("%d" % x)
            if y != self.lastY:
                self.yLabel.SetLabel("%d" % y)
        except TypeError, e:
                pass
        if c != self.lastC:
            self.cLabel.SetLabel("%06X" % c)
            
        #move frame if mouse is over it
        try:
            (mainscreenwidth, mainscreenheight) = wx.DisplaySize()
        except TypeError, e: #use previous value if DisplaySize returns an error,
                             #which it randomly does out of my understanding
            pass
        frameW = self.GetRect().GetWidth()
        frameH = self.GetRect().GetHeight()
        (frameX, frameY) = self.GetScreenPosition()
        if self.GetRect().ContainsXY(x, y):
            if frameX < mainscreenwidth/2:
                if frameY < mainscreenheight/2: #upper left; go to bottom left
                    self.Move((frameX, mainscreenheight - frameH))
                else: #bottom left; go to bottom right
                    self.Move((mainscreenwidth - frameW, mainscreenheight - frameH))
            else:
                if frameY < mainscreenheight/2: #upper right; go to upper left
                    self.Move((0, 0))
                else: #bottom right; go to upper right
                    self.Move((mainscreenwidth - frameW, 0))
            
        if (x != self.lastX) or (y != self.lastY):
        
            #repaint deskframe
            deskDC = win32ui.CreateDCFromHandle(win32gui.GetWindowDC(self.deskWnd))
            deskDC.BitBlt((0, 0), (self.screenwidth, self.screenheight),
                          self.desktopDC, (0, 0), win32con.SRCCOPY)
            
            if self.snapPoint1:
                self.drawSnapPoint(self.snapPoint1)
            
            #update magnified bitmap
            zoomDC = wx.WindowDC(self.magnifyBitmap)
            deskDC = self.getWXDCfromHandle(self.deskWnd)
            zoomDC.SetUserScale(8, 8)
            zoomDC.Blit(0, 0, 32, 32, deskDC, x-16, y-16)
            
            #draw an inverted colors crosshair on it
            zoomDC.SetLogicalFunction(wx.INVERT)
            zoomDC.SetBrush(wx.Brush(wx.Colour(0, 0, 0)))
            for cx in range (0, 32):
                zoomDC.DrawRectangle(cx, 16, 1, 1)
            for cy in range (0, 32):
                if cy != 16:
                    zoomDC.DrawRectangle(16, cy, 1, 1)            
        
        self.lastX = x
        self.lastY = y
        self.lastC = c

        
# end of class magnifierFrame

if __name__ == "__main__":
    def onSnap(event):
        print "Snap!"
        im = magnifierFrame.screenshotim
        im.save('screenshot.png', 'PNG')
        app.Exit()
        
    def onEsc(event):
        print "Cancelled"
        app.Exit()
        
    app = wx.PySimpleApp(0)
    wx.InitAllImageHandlers()
    magnifierFrame = magnifierFrame(app, True, None, -1, "")
    app.SetTopWindow(magnifierFrame)
    app.Bind(magnifierFrame.EVT_SNAP_EVENT, onSnap)  
    app.Bind(magnifierFrame.EVT_ESC_EVENT, onEsc)  
    magnifierFrame.Show()
    app.MainLoop()
